package com.sam.book.tree;

/**
 * Created by samwang on 2017/12/23.
 */
public class AvlTree<E extends Comparable<? super E>> {

    private AvlNode<E> root;

    public void insert(E element) {
        root = insert(element, root);
    }
    /**
     * 获取节点高度的方法
     * @param node
     * @return
     */
    private int getHeight(AvlNode<E> node) {
        return node == null ? -1 : node.height;
    }

    private int computeHeight(AvlNode<E> node) {
        if (node == null) {
            return -1;
        } else {
            node.height = Math.max(getHeight(node.left), getHeight(node.right));
            return node.height;
        }
    }

    public void remove(E e) {
        root = remove(e, root);
    }

    private AvlNode<E> findMin(AvlNode<E> node) {
        if (node.left != null) {
            return findMin(node.left);
        }
        return node;
    }

    /**
     * 移除后要将其子节点链到树中
     * @param element
     * @param node
     * @return
     */
    private AvlNode<E> remove(E element, AvlNode<E> node) {
        if (node == null) {
            return null;//不存在该元素，无需删除
        }

        int compare = element.compareTo(node.element);
        if (compare == 0) {
            if (node.left != null && node.right != null) {
                //被移除的节点左右都有节点，右子节点代替当前节点，左子节点放在右子节点最小节点的左节点
                findMin(node.right).left = node.left;
                node = node.right;
            } else {
                node = node.left == null ? node.right : node.left;
            }

            return node;
        } else if (compare > 0) {
            node.right = remove(element, node.right);
        } else {
            node.left = remove(element, node.left);
        }

        return balance(node);
    }

    private AvlNode<E> insert(E element, AvlNode<E> node) {
        if (node == null) {
            return new AvlNode<E>(element);
        } else {
            int compareResult = element.compareTo(node.element);
            if (compareResult == 0) {
                //重复值插入不处理
            } else if (compareResult > 0) {
                node.right = insert(element, node.right);
            } else {
                node.left = insert(element, node.left);
            }
        }
        return node;
    }

    /**
     * 平衡
     * @param node
     * @return
     */
    private AvlNode<E> balance(AvlNode<E> node) {
        if (node == null) {
            return null;
        }

        if (getHeight(node.left) - getHeight(node.right) > 2) {//左树高于右树大于2
            if (getHeight(node.left.left) >= getHeight(node.left.right)) {
                //左-左
                node = rotateWithLeftChild(node);
            } else {
                //左-右
                node = doubleWithLeftChild(node);
            }

        } else if (getHeight(node.right) - getHeight(node.left) > 2) {
            if (getHeight(node.left.left) >= getHeight(node.left.right)) {
                //右-左
                node = doubleWithRightChild(node);
            } else {
                //右-右
                node = rotateWithRightChild(node);
            }
        }

        return node;
    }

    /**
     * 第一种情况。左-左，node-k2,left-k1
     * @param node
     * @return
     */
    private AvlNode<E> rotateWithLeftChild(AvlNode<E> node) {
        AvlNode<E> left = node.left;
        node.left = left.right;
        left.right = node;
        computeHeight(node);
        computeHeight(left);
        return left;
    }
    /**
     * 第四种情况右单旋转。右-右
     * @param node
     * @return
     */
    private AvlNode<E> rotateWithRightChild(AvlNode<E> node) {
        AvlNode<E> right = node.right;
        node.right = right.left;
        right.right = node;
        computeHeight(node);
        computeHeight(right);
        return right;
    }

    private AvlNode<E> doubleWithLeftChild(AvlNode<E> node) {
        node.left = rotateWithRightChild(node.left);
        return rotateWithLeftChild(node);
    }

    private AvlNode<E> doubleWithRightChild(AvlNode<E> node) {
        node.right = rotateWithLeftChild(node.right);
        return rotateWithRightChild(node);
    }






    private static class AvlNode<T>{
        public AvlNode(T element) {
            this.element = element;
        }

        T element;
        AvlNode<T> left;
        AvlNode<T> right;
        int height;

    }

}
